import joblib
import gradio as gr
import numpy as np
from PIL import Image

# 加载保存的KNN模型
best_knn_model = joblib.load('best_knn_model.pkl')


def predict_digit(drawing):
    try:
        # 获取PIL 图像
        image = drawing['composite']

        # 调整图像大小为 8x8
        image = image.resize((8, 8), Image.Resampling.LANCZOS)

        # 将 PIL 图像转换为 NumPy 数组
        image_array1 = np.array(image)

        image_array2 = image_array1[:, :, 3]
        threshold1 = 30
        threshold2 = 125
        image_array2[image_array2 < threshold1] = 0
        image_array2[image_array2 > threshold2] = 125
        image_array2 = image_array2 * 2
        image_orgin = Image.fromarray(image_array2)

        # 查看数组形状
        # print(image_array2)

        # 显示图像
        # image_orgin.show()

        # 展平图像
        image_array3 = image_array2.reshape((1, -1))

        # 进行预测并返回结果
        prediction = best_knn_model.predict(image_array3)[0]
        return str(prediction)
    except Exception as e:
        return f"Error: {e}"


def create_gradio_interface():
    inp = gr.Sketchpad(label="Draw a digit", type='pil')  # 创建 Sketchpad 组件并指定类型
    iface = gr.Interface(fn=predict_digit,
                         inputs=inp,
                         outputs="text",
                         live=False)

    return iface


if __name__ == "__main__":
    demo = create_gradio_interface()
    demo.launch()
